package com.yys.apm.collect;


import com.yys.apm.boot.AgentBootLoader;
import com.yys.apm.model.CollectInfo;
import com.yys.apm.boot.spi.CollectSPI;
import com.yys.apm.constant.CommonConstant;
import com.yys.apm.utils.ApmAgentUtil;
import javassist.CtMethod;
import org.apache.ibatis.executor.statement.BaseStatementHandler;

import javax.servlet.http.HttpServletRequest;
import java.sql.Connection;
import java.sql.SQLException;

/**
 * Mybatis采集器
 *
 * @author Yingsheng Ye
 */
@CollectSPI
public class MybatisCollect extends AbstractApmCollects implements ApmCollect {

    public static MybatisCollect INSTANCE;

    public MybatisCollect() {
        INSTANCE = this;
    }

    private static final String startSrc;
    private static final String endSrc;
    private static final String errorSrc;

    static {
        StringBuilder sbuilder = new StringBuilder();
        sbuilder.append("com.yys.apm.collect.MybatisCollect instance= ");
        sbuilder.append("com.yys.apm.collect.MybatisCollect.INSTANCE;\r\n");
        sbuilder.append("com.yys.apm.model.CollectInfo collectInfo=instance.start($0,$1,\"%s\",\"%s\");");
        startSrc = sbuilder.toString();
        sbuilder = new StringBuilder();
        sbuilder.append("instance.end(collectInfo);");
        endSrc = sbuilder.toString();
        sbuilder = new StringBuilder();
        sbuilder.append("instance.error(collectInfo,e);");
        errorSrc = sbuilder.toString();
    }

    public CollectInfo start(BaseStatementHandler handler, Connection connection, String className, String methodName) throws SQLException {
        String sql = handler.getBoundSql().getSql();
        String params = ApmAgentUtil.objectToJson(handler.getParameterHandler().getParameterObject());
        HttpServletRequest httpServletRequest = ApmAgentUtil.getHttpServletRequest();
        CollectInfo collectInfo = CollectInfo.build().builder(httpServletRequest, MybatisCollect.class.getSimpleName(),
                null != httpServletRequest ? (String) httpServletRequest.getAttribute(CommonConstant.TRACE_ID) : null,
                null != httpServletRequest ? (String) httpServletRequest.getAttribute(CommonConstant.SPAN_ID) : null,false);
        collectInfo.setClassName(className);
        collectInfo.setMethodName(methodName);
        collectInfo.put("sql", sql);
        collectInfo.put("params", params);
        collectInfo.put("jdbcUrl", connection.getMetaData().getURL());
        return collectInfo;
    }

    public void error(CollectInfo collectInfo, Throwable e) {
        collectInfo.setStatus(CommonConstant.FAILED);
        collectInfo.setMsg(e.getMessage());
        collectInfo.put("stackTrace",e.getStackTrace());
    }

    public void end(CollectInfo collectInfo) {
        collectInfo.setEndTime(System.currentTimeMillis());
        collectInfo.setCostTime(collectInfo.getEndTime() - collectInfo.getStartTime());
        this.execute(collectInfo);
    }

    @Override
    public byte[] transform(AgentBootLoader agentLoader) throws Exception {
        CtMethod connectMethod = agentLoader.getCtClass().getMethod("prepare", "(Ljava/sql/Connection;Ljava/lang/Integer;)Ljava/sql/Statement;");
        agentLoader.agentMethod(connectMethod, setBuild(agentLoader.getClassName(), connectMethod));
        return agentLoader.getCtClass().toBytecode();
    }

    @Override
    public boolean isApmTarget(AgentBootLoader agentLoader) {
        return agentLoader.getClassName().equals("org.apache.ibatis.executor.statement.BaseStatementHandler");
    }


    @Override
    protected AgentBootLoader.MethodSourceBuild setBuild(String className, CtMethod method) {
        AgentBootLoader.MethodSourceBuild build = new AgentBootLoader.MethodSourceBuild();
        build.setStartSrc(String.format(startSrc, className, method.getName()));
        build.setErrorSrc(errorSrc);
        build.setEndSrc(endSrc);
        return build;
    }

}
